import Foundation
import Accelerate

public struct Algorithm {
    var sampleRate: Int
    public init(sampleRate: Int) {
        self.sampleRate = sampleRate
    }
    
    public func dft(RedBuffer: [Double], IRBufer: [Double]) -> Double {
        var R: Double
        var ACRed, DCRed, ACIR, DCIR: Double
        DCIR = vDSP.sum(IRBufer)
        DCRed = vDSP.sum(RedBuffer)
        let dftType = try? vDSP.DiscreteFourierTransform(previous: nil, count: RedBuffer.count, direction: .forward, transformType: .complexReal, ofType: Double.self)
        var dftResult = dftType?.transform(real: RedBuffer, imaginary: [Double](repeating: 0.0, count: RedBuffer.count))
        ACRed = dftResult!.real[1] * dftResult!.real[1] + dftResult!.imaginary[1] * dftResult!.imaginary[1]
        dftResult = dftType?.transform(real: IRBufer, imaginary: [Double](repeating: 0.0, count: IRBufer.count))
        ACIR = dftResult!.real[1] * dftResult!.real[1] + dftResult!.imaginary[1] * dftResult!.imaginary[1]
        R = ACRed / DCRed / ACIR * DCIR
        return R
    }
    
    public func getSPO2(R: Double) -> Double {
        return -45.060 * R * R + 30.354 * R + 94.845
    }
    
    public func getR(IRdata: [Double], redData: [Double]) -> Double {
        let cosinePart = Double(cos(-2.0 * Double.pi / Double(IRdata.count) * 2.0))
        let sinePart = Double(sin(-2.0 * Double.pi / Double(IRdata.count) * 2.0))
        let DCRed = vDSP.sum(redData)
        let DCIR = vDSP.sum(IRdata)
        var ACRed: Double = 0.0
        var ACIR: Double = 0.0
        var ACRed_real: Double = 0.0
        var ACRed_imag: Double = 0.0
        var ACIR_real = 0.0
        var ACIR_imag: Double = 0.0
        for i in 0..<redData.count {
            ACIR_real += IRdata[i] * cosinePart * Double(i)
            ACIR_imag += IRdata[i] * sinePart * Double(i)
            ACRed_real += redData[i] * cosinePart * Double(i)
            ACRed_imag += redData[i] * sinePart * Double(i)
        }
        ACIR = sqrt(Double(ACIR_real * ACIR_real + ACIR_imag * ACIR_imag))
        ACRed = sqrt(ACRed_real * ACRed_real + ACRed_imag * ACRed_imag)
        let R = ACRed / DCRed / ACIR * DCIR
        return R
    }
    
    public func dft2(input: [Double]) -> [Double] {
        let len = input.count
        let real = [Double](input)
        let imag = [Double](repeating: 0.0, count: len)
        let dftType = try? vDSP.DiscreteFourierTransform.init(count: len, direction: .forward, transformType: .complexComplex, ofType: Double.self)
        let complexOutput = dftType?.transform(real: real, imaginary: imag)
        guard var clxOpt = complexOutput else {
            print("Empty Pointer error")
            return [0.0, 0.0]
        }
        var spl = DSPDoubleSplitComplex(realp: &clxOpt.real, imagp: &clxOpt.imaginary)
        var sqrt = [Double](repeating: 0.0, count: len)
        vDSP.absolute(spl, result: &sqrt)
        return sqrt
    }
    
    func cutDC(data: [Double]) -> Int {
        var index = 0
        for i in 3..<data.count {
            if (data[i] > data[i - 1]) && (data[i] > data[i + 1]) {
                index = i
                break
            }
        }
        return index
    }
    
    public func getHRSPO2(IRBuffer: [Double], RedBuffer: [Double]) ->[Int16] {
        let DFT_IR = dft2(input: IRBuffer)
        let DFT_Red = dft2(input: RedBuffer)
        let AC_IR = DFT_IR.suffix(from: 3).prefix(IRBuffer.count / 2).max()
        let AC_Red = DFT_Red.suffix(from: 3).prefix(RedBuffer.count / 2).max()
        guard let AC_IR2 = AC_IR, let AC_Red2 = AC_Red else {
            print("NO MAX")
            return [-1, -1]
        }
        let R = AC_Red2 / DFT_Red[0] / AC_IR2 * DFT_IR[0]
        let F0 = Double(self.sampleRate) / Double(IRBuffer.count)
        let index = cutDC(data: DFT_IR)
        let heartRate = Int16(Double(index * 60) * F0)
        let SPO2 = Int16(getSPO2(R: R))
        return [heartRate, SPO2]
    }
}
